import PySimpleGUI as sg
import shutil
from Crypto.Hash import SHA
from Crypto.Cipher import ARC4
from struct import unpack
from tempfile import NamedTemporaryFile
from pathlib import Path
from __version__ import __version__


def decryptSQLite(in_file, out_file, password):

    ret = None

    with open(in_file, 'rb') as f:
        key = SHA.new(password).digest()[:16]
        header = ARC4.new(key).decrypt(f.read(1024))
        if header[0:15] == b'SQLite format 3':
            declared_ps = unpack('>H', header[16:18])[0]
            if declared_ps == 1:
                declared_ps = 65536
            t = NamedTemporaryFile(delete=False, suffix='.sqlite')
            f.seek(0)

            while True:
                block = f.read(declared_ps)
                if not block:
                    break
                t.write(ARC4.new(key).decrypt(block))  # reinitialize rc4
            t.close()
            ret = t.name
    if ret:
        shutil.move(ret, out_file)
    return ret


def main():

    font_yh = ('微软雅黑', 10)

    layout = [
        [
            sg.Text('目标文件：', font=font_yh),
            sg.Input(enable_events=True, disabled=True),
            sg.FileBrowse('选择文件', key='FILE', font=font_yh, size=(8, 1)),
        ],
        [
            sg.Text('解密密码：', font=font_yh),
            sg.InputText(password_char='*', key='PASSWD', default_text='abcde'),
            sg.Button('解密', font=font_yh, size=(8, 1), key='DECRYPT')
        ],
        # [
        #     sg.Column([[sg.Button('解密', font=font_yh, size=(8, 1), key='DECRYPT')]], justification='center')
        # ]
    ]

    window = sg.Window(f'SQLite DB解密工具 v{__version__}', layout, size=(600, 120))

    # 事件循环
    while True:
        event, values = window.read()

        if event == sg.WIN_CLOSED:
            break

        if event == 'DECRYPT':
            src_path = values['FILE']
            des_path = Path(src_path).parent / f"{Path(src_path).stem}_bak{Path(src_path).suffix}"
            passwd = values["PASSWD"].encode()

            # 校验
            if not src_path and not passwd:
                sg.popup_notify('目标文件和密码不能为空！')
                continue
            if not src_path:
                sg.popup('请先选择要解密的文件！', title='', font=font_yh)
            if not passwd:
                sg.popup('请输入解密密码！', title='', font=font_yh)
                continue

            # 执行解密
            res = decryptSQLite(src_path, des_path, passwd)
            if res:
                sg.popup_notify('解密文件已保存至源文件所在目录！', title='解密成功')
            else:
                sg.popup(f'解密失败！', font=font_yh)

    window.close()


if __name__ == '__main__':
    main()
